"""
2.	机器人是模仿人类和动物行为的机器，其内部有台计算机，通过读取各个传感器的信息，做出判断，并且调用电机实现相关的动作，完成指令。
给定“手势识别”数据集，有三个指令：左转、右转、停止。利用深度学习平台预训练模型mobilenet，搭建后端网络，进行模型训练和测试，按下面的要求，
完成相应代码（25分）
"""
import sys
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, optimizers, losses, metrics
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from python_ai.common.xcommon import *
from python_ai.common.read_data.redis_numpy import fromRedisNd
from python_ai.cate.redis.vm_ubun20_redis import r

SIDE = 224
ALPHA = 1e-3
BATCH_SIZE = 32
N_EPOCH = 10
BASE_DIR, FILE_NAME = os.path.split(__file__)

# ③	创建模型：包括预训练模型mobilenet和后端网络，实现3种手势分类任务
inputs = keras.Input((SIDE, SIDE, 3))
base_model = keras.applications.mobilenet.MobileNet(
    include_top=False,
    weights='imagenet',
    pooling='avg'
)
base_model.trainable = False
customer_model = layers.Dense(3, activation=activations.softmax)
x = base_model(inputs)
x = customer_model(x)
model = keras.Model(inputs, x)
model.summary()

# ①	导入“手势识别”数据集
x, y, path = [], [], []
idx2label, label2idx = {}, {}
print('Loading from redis...')
r_key = ensure_filename(BASE_DIR, True) + 'v1'
print('redis key', r_key)
dt1 = datetime.datetime.now()
x = fromRedisNd(r, f'{r_key}_x')
y = fromRedisNd(r, f'{r_key}_y')
path = fromRedisNd(r, f'{r_key}_path')
labels_arr = fromRedisNd(r, f'{r_key}_labels_arr')
dt2 = datetime.datetime.now()
print(f'Loaded. Time usage: {dt2 - dt1}')
idx2label = {idx: label for idx, label in enumerate(labels_arr)}
label2idx = {label: idx for idx, label in enumerate(labels_arr)}
print('x', x.shape)
print('y', y.shape)
print('path', np.shape(path))
print('labels_arr', labels_arr)
print('idx2label', idx2label)
print('label2idx', label2idx)

# ②	按适当比例划分训练集、验证集、测试集
x_train, x_val_test, y_train, y_val_test, _, path_val_test = train_test_split(x, y, path, train_size=0.8, random_state=1, shuffle=True)
x_val, x_test, y_val, y_test, _, path_test = train_test_split(x_val_test, y_val_test, path_val_test, train_size=0.5, random_state=1, shuffle=True)

# ④	进行模型编译和训练，打印输出训练集、验证集、测试集准确率
model.compile(
    optimizer=optimizers.Adam(learning_rate=ALPHA),
    loss=losses.sparse_categorical_crossentropy,
    metrics=[metrics.sparse_categorical_accuracy]
)
model.fit(
    x_train,
    y_train,
    batch_size=BATCH_SIZE,
    epochs=N_EPOCH,
    validation_data=(x_val, y_val)
)
print('Testing ...')
model.evaluate(x_test, y_test, BATCH_SIZE)
print('Tested')

# ⑤	从测试集中随机选取9张图片，利用训练模型，图示预测种类名称、并与真实种类名称比较，预测正确显示“黑色”，预测错误显示“红色”。
pred = model.predict(x_test[:9]).argmax(axis=1)
spr = 3
spc = 3
spn = 0
plt.figure(figsize=[7, 7])
for i in range(spr * spc):
    real_idx = y_test[i]
    pred_idx = pred[i]
    right = real_idx == pred_idx
    path = path_test[i]
    img = cv.imread(path, cv.IMREAD_COLOR)[:, :, ::-1]
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.imshow(img)
    plt.axis('off')
    plt.title(f'{idx2label[real_idx]}=>{idx2label[pred_idx]}', color='black' if right else 'red')

plt.show()
